import cv2
import os
import numpy as np
import argparse
from facenet_pytorch import MTCNN

class OpenCVFaceRecognitionSystem:
    def __init__(self, display_mode=True, debug_mode=False, camera_id=0, 
                 display_target='local', confidence_threshold=100):
        self.display_mode = display_mode
        self.debug_mode = debug_mode
        self.camera_id = camera_id
        self.display_target = display_target  # 'local' 或 'remote'
        self.confidence_threshold = confidence_threshold
        
        # 创建LBPH人脸识别器
        self.face_recognizer = cv2.face.LBPHFaceRecognizer_create()
        self.known_face_names = []
        self.known_face_ids = {}
        self.next_id = 0
        
        # 加载MTCNN人脸检测器
        self.face_detector = MTCNN(keep_all=True, device='cpu')
        
        # 加载已知人脸
        self.load_known_faces("known_faces")
        
    def load_known_faces(self, folder_path):
        """加载已知人脸"""
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)
            print(f"请在 {folder_path} 文件夹中添加已知人脸图像")
            return
            
        # 收集训练数据
        faces = []
        labels = []
        
        for file in os.listdir(folder_path):
            if file.endswith((".jpg", ".png", ".jpeg")):
                image_path = os.path.join(folder_path, file)
                # 读取彩色图像
                image = cv2.imread(image_path)
                
                if image is None:
                    print(f"警告: 无法读取图像 {file}")
                    continue
                    
                # 使用MTCNN检测人脸
                # facenet-pytorch的MTCNN返回不同格式的结果
                boxes, _ = self.face_detector.detect(image)
                detections = []
                if boxes is not None:
                    for box in boxes:
                        x, y, x2, y2 = box
                        width, height = x2 - x, y2 - y
                        detections.append({
                            'box': [int(x), int(y), int(width), int(height)],
                            'confidence': 1.0  # facenet-pytorch不直接返回置信度
                        })
                
                if len(detections) > 0:
                    # 取第一个检测到的人脸
                    face_data = detections[0]
                    x, y, width, height = face_data['box']
                    # 确保坐标不为负数
                    x, y = max(0, x), max(0, y)
                    
                    # 提取人脸区域（转换为灰度图）
                    face_roi = image[y:y+height, x:x+width]
                    face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
                    face_roi = cv2.resize(face_roi, (100, 100))
                    
                    # 从文件名中提取人名作为标签（去除数字部分）
                    name = file.split(".")[0]
                    # 移除文件名中的数字部分，只保留人名
                    name = ''.join([char for char in name if not char.isdigit()])
                    # 处理可能的空格
                    name = name.strip()
                    
                    if name not in self.known_face_ids:
                        self.known_face_ids[name] = self.next_id
                        self.next_id += 1
                        
                    faces.append(face_roi)
                    labels.append(self.known_face_ids[name])
                    self.known_face_names.append(name)
                    print(f"已加载: {file} (标签: {name})")
                else:
                    print(f"警告: 在 {file} 中未检测到人脸")
        
        # 训练识别器
        if len(faces) > 0 and len(labels) > 0:
            self.face_recognizer.train(faces, np.array(labels))
            print(f"训练完成，共加载 {len(faces)} 个人脸样本")
        else:
            print("未找到有效的人脸样本进行训练")
    
    def recognize_faces(self, frame):
        """识别人脸"""
        # 使用MTCNN检测人脸
        # facenet-pytorch的MTCNN返回不同格式的结果
        boxes, _ = self.face_detector.detect(frame)
        detections = []
        if boxes is not None:
            for box in boxes:
                x, y, x2, y2 = box
                width, height = x2 - x, y2 - y
                detections.append({
                    'box': [int(x), int(y), int(width), int(height)]
                })
        
        face_names = []
        face_locations = []
        
        for detection in detections:
            x, y, width, height = detection['box']
            # 确保坐标不为负数
            x, y = max(0, x), max(0, y)
            
            face_locations.append((y, x+width, y+height, x))  # (top, right, bottom, left) 原始坐标格式
            
            # 提取人脸区域（转换为灰度图）
            face_roi = frame[y:y+height, x:x+width]
            face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
            face_roi = cv2.resize(face_roi, (100, 100))
            
            # 识别人脸
            if len(self.known_face_ids) > 0:
                label, confidence = self.face_recognizer.predict(face_roi)
                
                if self.debug_mode:
                    print(f"识别结果: 标签={label}, 置信度={confidence}")
                
                # 根据置信度判断是否匹配
                if confidence < self.confidence_threshold:
                    # 通过标签找到名称
                    name = None
                    for n, l in self.known_face_ids.items():
                        if l == label:
                            name = n
                            break
                    if name:
                        face_names.append(name)
                    else:
                        face_names.append("未知")
                else:
                    face_names.append(f"未知({confidence:.1f})")
            else:
                face_names.append("人脸(无匹配)")
        
        return face_locations, face_names
    
    def run(self):
        """运行人脸识别系统"""
        video_capture = cv2.VideoCapture(self.camera_id)
        
        # 检查摄像头是否成功打开
        if not video_capture.isOpened():
            print(f"错误: 无法打开摄像头 {self.camera_id}")
            # 尝试其他常见的设备ID
            for i in range(1, 5):
                print(f"尝试摄像头 {i}...")
                video_capture = cv2.VideoCapture(i)
                if video_capture.isOpened():
                    print(f"成功打开摄像头 {i}")
                    self.camera_id = i
                    break
            if not video_capture.isOpened():
                return
        
        # 设置分辨率
        video_capture.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        video_capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        
        if self.display_mode:
            print("OpenCV人脸识别系统已启动，按 'q' 退出")
        else:
            print("OpenCV人脸识别系统已启动（无GUI模式），按 Ctrl+C 退出")
        
        frame_count = 0
        try:
            while True:
                ret, frame = video_capture.read()
                
                if not ret:
                    print("无法获取摄像头画面")
                    break
                
                frame_count += 1
                
                # 识别人脸
                face_locations, face_names = self.recognize_faces(frame)
                
                # 打印识别结果到控制台
                if len(face_names) > 0:
                    print(f"检测到人脸: {', '.join(face_names)}")
                elif self.debug_mode and frame_count % 30 == 0:
                    print("未检测到人脸")
                
                if self.display_mode:
                    # 显示结果
                    for (top, right, bottom, left), name in zip(face_locations, face_names):
                        # 画人脸框
                        # 注意：OpenCV的坐标顺序是(left, top, right, bottom)
                        cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)
                        
                        # 画标签背景
                        cv2.rectangle(frame, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED)
                        font = cv2.FONT_HERSHEY_DUPLEX
                        cv2.putText(frame, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1)
                    
                    # 显示图像
                    try:
                        cv2.imshow('OpenCV人脸识别系统', frame)
                    except cv2.error as e:
                        print(f"显示错误: {e}")
                        print("提示: 如果通过SSH连接，请使用 --no-gui 参数运行")
                        print("或者设置正确的显示环境变量")
                        # 继续运行但不显示
                        self.display_mode = False
                        continue
                    
                    # 按q退出
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        break
                else:
                    # 添加小延迟以控制CPU使用率
                    cv2.waitKey(100)
        
        except KeyboardInterrupt:
            print("\n用户中断程序")
        finally:
            # 释放资源
            video_capture.release()
            if self.display_mode:
                cv2.destroyAllWindows()

def test_camera():
    """测试摄像头是否正常工作"""
    print("测试摄像头...")
    cap = cv2.VideoCapture(0)
    
    if not cap.isOpened():
        print("错误: 无法打开摄像头")
        return
    
    print("摄像头已打开，按任意键退出测试")
    while True:
        ret, frame = cap.read()
        if not ret:
            print("无法获取画面")
            break
            
        # 显示简单的信息
        cv2.putText(frame, "摄像头测试 - 按任意键退出", (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        
        cv2.imshow('摄像头测试', frame)
        
        if cv2.waitKey(1) != -1:  # 按任意键退出
            break
    
    cap.release()
    cv2.destroyAllWindows()
    print("摄像头测试完成")

def main():
    parser = argparse.ArgumentParser(description='OpenCV人脸识别系统')
    parser.add_argument('--no-gui', action='store_true', help='无GUI模式运行')
    parser.add_argument('--debug', action='store_true', help='调试模式')
    parser.add_argument('--test-camera', action='store_true', help='测试摄像头')
    parser.add_argument('--camera-id', type=int, default=0, help='摄像头设备ID (默认: 0)')
    parser.add_argument('--display', choices=['local', 'remote'], default='local', 
                        help='显示目标: local(本地) 或 remote(远程) (默认: local)')
    parser.add_argument('--confidence', type=float, default=100.0, 
                        help='置信度阈值 (默认: 100.0)')
    args = parser.parse_args()
    
    if args.test_camera:
        test_camera()
        return
    
    try:
        face_recognition_system = OpenCVFaceRecognitionSystem(
            display_mode=not args.no_gui, 
            debug_mode=args.debug,
            camera_id=args.camera_id,
            display_target=args.display,
            confidence_threshold=args.confidence
        )
        face_recognition_system.run()
    except cv2.error as e:
        if "face" in str(e).lower() and "lbph" in str(e).lower():
            print("错误: OpenCV的face模块未安装")
            print("请运行以下命令安装:")
            print("pip install opencv-contrib-python")
        else:
            print(f"OpenCV错误: {e}")
    except Exception as e:
        # 检查是否是MTCNN相关错误
        if "mtcnn" in str(e).lower() or "facenet" in str(e).lower():
            print("错误: MTCNN库未安装")
            print("请运行以下命令:")
            print("pip install facenet-pytorch")
        else:
            print(f"发生错误: {e}")

if __name__ == "__main__":
    main()